####################################################
# INPUT: frames folders
# OUTPUT: landmarks of each frames
####################################################

import cv2
import os, sys
import glob
import dlib
import numpy as np
import time
import pdb
import argparse
from pathlib2 import Path

detector = dlib.get_frontal_face_detector()
predictor = dlib.shape_predictor('../data/face/shape_predictor_68_face_landmarks.dat')

def rect_to_bbox(rect):
    x = rect.left()
    y = rect.top()
    w = rect.right() - x
    h = rect.bottom() - y

    return (x, y, w, h)


def shape_to_np(landmk, dtype="int"):
    # initialize the list of (x, y)-coordinates
    coords = np.zeros((landmk.num_parts, 2), dtype=dtype)

    # loop over all facial landmarks and convert them
    # to a 2-tuple of (x, y)-coordinates
    for i in range(0, landmk.num_parts):
        coords[i] = (landmk.part(i).x, landmk.part(i).y)

    # return the list of (x, y)-coordinates
    return coords

def detect_image(path_image, path_landmk, path_landmk_5p, path_landmk_bbox, image=None, is_save=True):

    # path_save = path_image.parent.joinpath(path_image.stem + ".txt")
    if not path_image.exists():
        raise(f"ERROR: {path_image} not exist")

    if image is None:
        image = cv2.imread(str(path_image))

    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    rects = detector(gray, 1)

    if len(rects) == 0:
        return False
    elif len(rects) > 1:  # maybe have multi person
        print(f"WARRNING: have {len(rects)} face in this images!!!! select the first one!!!")
    
    for (i, rect) in enumerate(rects):
        if i>0: break

        landmk = predictor(gray, rect)
        landmk = shape_to_np(landmk)
        landmk_bbox = rect_to_bbox(rect)

        #print(f"rect type = {type(rect)}, value={rect}, (x,y,w,h)={landmk_bbox}")

        landmk_5p = np.zeros((5, 2), dtype=np.int)        # 5 keypoints from landmark
        landmk_5p[0] = np.round(np.mean(landmk[36:42,:], axis=0)).astype("int")  # eye left
        landmk_5p[1] = np.round(np.mean(landmk[42:48,:], axis=0)).astype("int")  # eye right 
        landmk_5p[2] = landmk[33]                                                # nose
        landmk_5p[3] = landmk[48]                                                # mouth left
        landmk_5p[4] = landmk[54]                                                # mouth right

        if is_save :
            np.savetxt(str(path_landmk), landmk,    fmt="%d", delimiter=" ", newline="\n", encoding="utf-8")
            np.savetxt(str(path_landmk_5p), landmk_5p, fmt="%d", delimiter=" ", newline="\n", encoding="utf-8")
            np.savetxt(str(path_landmk_bbox), landmk_bbox, fmt="%d", delimiter=" ", newline="\n", encoding="utf-8")
            
            
        # if show_image:   #for show the image
        #   for (x, y) in landmk:
        #       cv2.circle(image, (x, y), 1, (0, 0, 255), -1)
    
    return True


def arg_parser():
    parser = argparse.ArgumentParser("Extract video frames")
    parser.add_argument("--path_frames", type=str, default=None, help="video path for extract")
    parser.add_argument("--path_save", type=str, default="./landmk", help='save extraction frames')
    parser.add_argument("--override", action='store_true', help="if override origin frames")
    args = parser.parse_args()
    return args

if __name__ == "__main__":

    args = arg_parser()

    # create data path
    path_save = Path(args.path_save)

    if args.path_frames is None:
        path_frames = path_save.joinpath("frames")
    else:
        path_frames = Path(args.path_frames)

    path_landmk = path_save.joinpath("landmk")
    path_landmk_5p = path_save.joinpath("landmk_5p")
    path_landmk_bbox = path_save.joinpath("landmk_bbox")

    path_landmk.mkdir(parents=True, exist_ok=True)
    path_landmk_5p.mkdir(parents=True, exist_ok=True)
    path_landmk_bbox.mkdir(parents=True, exist_ok=True)
    
    # STEP1: get frames
    v_path_frames = sorted(path_frames.glob("*.jpg"))
    n_frame = len(v_path_frames)
    print(f"have {n_frame} frames in {path_frames}")

    # STEP2: extract landmarks
    for i, ipath_frame in enumerate(v_path_frames):
        frame_id = int(ipath_frame.stem[-4:])

        # detect image
        image = cv2.imread(str(ipath_frame)) 

        
        ipath_landmk_frame = path_landmk.joinpath(f"frame-{frame_id:04d}.txt")   
        ipath_landmk_frame_5p = path_landmk_5p.joinpath(f"frame-{frame_id:04d}.txt")    
        ipath_landmk_frame_bbox = path_landmk_bbox.joinpath(f"frame-{frame_id:04d}.txt")


        if args.override or not ipath_landmk_frame.exists() or \
            not ipath_landmk_frame_5p.exists() or not ipath_landmk_frame_bbox.exists():

            detect_image(ipath_frame, 
                        path_landmk=ipath_landmk_frame, 
                        path_landmk_5p=ipath_landmk_frame_5p, 
                        path_landmk_bbox=ipath_landmk_frame_bbox,
                        image=image, is_save=True)

        if frame_id % 50 == 0:
            print(f"extract and save the {frame_id}/{n_frame}-th frame --> {ipath_landmk_frame}")
